Using YOLO v3¶
Notebook to test a pre-trained YOLO on the COCO dataset and trying to apply to aerial images.
This notebook was made based on the tutorial avaiable on https://www.pyimagesearch.com/2018/11/12/yolo-object-detection-with-opencv/
In [66]:
# import the necessary packages
import numpy as np
import os
import cv2
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (20.0, 15.0)
%matplotlib notebook
Read Classes¶
In [2]:
# read the COCO class labels
labels_path = "./yolo-coco/coco.names"
LABELS = open(labels_path).read().strip().split("\n")
In [3]:
len(LABELS)
Out[3]:
In [4]:
# list of colors to represent each class
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(len(LABELS), 3),
dtype="uint8")
Import Pre-Trained YOLO¶
In [5]:
# read weights and config from pre-trained YOLO
weights_path = "./yolo-coco/yolov3.weights"
config_path = "./yolo-coco/yolov3.cfg"
nn_model = cv2.dnn.readNetFromDarknet(config_path, weights_path)
In [6]:
nn_model.getLayerNames()
Out[6]:
Function to detect objects in image¶
In [98]:
def detect_objects_yolo(nn_model, image_path, threshold_boxes = 0.3, min_probability = 0.5):
'''
Function to apply a forward pass on YOLO NN and apply Non-Max Supression on the detected bounding boxes
@param nn_model: trained YOLO model
@param image_path: path string to download the image
@param threshold_boxes: threshold to apply on non-max suppression
@param min_probability: minimum probability to filter class and bounding boxes detections
'''
# Load Image
image = cv2.imread(image_path)
# Get Output Layers Names
ln = nn_model.getUnconnectedOutLayersNames()
# Create Blob image and set as input (resizing image to 416x416).
# Then, run the forward to identify classes and bounding boxes
blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416),
swapRB=True, crop=False)
nn_model.setInput(blob)
output_layers = nn_model.forward(ln)
# For each output layer we have to identify each detection and filter them based on min_probability
bounding_boxes = []
class_confidences = []
class_prediction_ids = []
for output in output_layers:
for detection in output:
# First 5 params are for bounding box prediction
# [center_x, center_y, width, height, bb_confidence, (classes_confidence)]
class_scores = detection[5:]
class_prediction_id = np.argmax(class_scores)
class_confidence = class_scores[class_prediction_id]
if detection[4] >= min_probability:
# Use images size to find correct coordinates to draw the bounding boxes
(image_height, image_width) = image.shape[:2]
box = detection[0:4] * np.array([image_width, image_height, image_width, image_height])
(centerX, centerY, width, height) = box.astype("int")
# To draw the bounding boxes using OpenCV we have to get the top left coordinates
x = int(centerX - (width / 2))
y = int(centerY - (height / 2))
# Store values
bounding_boxes.append([x, y, int(width), int(height)])
class_confidences.append(float(class_confidence))
class_prediction_ids.append(class_prediction_id)
## Apply Non-Max Supression in the bounding boxes
indices = cv2.dnn.NMSBoxes(bounding_boxes, class_confidences, min_probability, threshold_boxes)
if len(indices) > 0:
for i in indices.flatten():
(x, y, w, h) = (bounding_boxes[i][0], bounding_boxes[i][1], bounding_boxes[i][2], bounding_boxes[i][3])
# Draw the detection on the image
color = [int(c) for c in COLORS[class_prediction_ids[i]]]
cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
class_name = "{}: {:.4f}".format(LABELS[class_prediction_ids[i]], class_confidences[i])
cv2.putText(image, class_name, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
fig=plt.figure(image_path)
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB));
In [99]:
images_path = './images/'
for file in os.listdir(images_path):
if file.endswith('.jpg') or file.endswith('.jpeg'):
detect_objects_yolo(nn_model, images_path+file, threshold_boxes = 0.3, min_probability = 0.5)
In [ ]: